import pybullet as p
import pybullet_data
import time
import numpy as np

# 连接到物理服务器
physicsClient = p.connect(p.GUI)
p.setGravity(0, 0, -9.8)

# 创建地面
p.createCollisionShape(p.GEOM_PLANE)
p.createMultiBody(0, 0)

# 添加数据路径并加载URDF
# p.setAdditionalSearchPath(pybullet_data.getDataPath())
# 假设您的URDF文件在当前目录下
robot_id = p.loadURDF("./tests/urdf/revolute_joint.urdf", [0, 0, 0.05])

# 获取关节信息
num_joints = p.getNumJoints(robot_id)
print(f"Number of joints: {num_joints}")

for i in range(num_joints):
    joint_info = p.getJointInfo(robot_id, i)
    print(f"Joint {i}: {joint_info[1].decode('utf-8')}, Type: {joint_info[2]}")

# 设置关节电机控制（位置控制）
target_position = np.pi / 4  # 45度
max_force = 500
position_gain = 0.5

p.setJointMotorControl2(
    bodyUniqueId=robot_id,
    jointIndex=0,  # 第一个关节
    controlMode=p.POSITION_CONTROL,
    targetPosition=target_position,
    force=max_force,
    positionGain=position_gain
)

# 仿真循环
try:
    for _ in range(100000):
        p.stepSimulation()
        
        # 获取当前关节状态
        joint_state = p.getJointState(robot_id, 0)
        print(f"Joint position: {joint_state[0]:.2f} rad")
        
        time.sleep(1./240.)
except KeyboardInterrupt:
    pass

p.disconnect()